package com.xbai.spark.sql

import org.apache.spark.SparkConf
import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql._

/**
  * 用户自定义聚合函数
  * 强类型的 Dataset 和弱类型的 DataFrame 都提供了相关的聚合函数，
  * 如 count()，countDistinct()，avg()，max()，min()。除此之外，用户可以设定自己的自定义聚合函数。
 *
  * 强类型用户自定义聚合函数：通过继承 Aggregator 来实现强类型自定义聚合函数，同样是求平均工资
  * 案例：求平均年龄的自定义聚合函数。
  *
  * @author xbai
  * @Date 2021/1/5
  */
object SparkSQL05_UDAF_Class {

  def main(args: Array[String]): Unit = {
    // 创建配置对象
    val conf: SparkConf = new SparkConf().setMaster("local[*]").setAppName("udaf")
    // 创建 SparkSQL 的环境对象
    val spark: SparkSession = SparkSession.builder().config(conf).getOrCreate()

    // 在进行转换之前，需要引入隐式转换规则
    import spark.implicits._

    val udaf = new MyAgeAvgClassFunction

    val avgCol: TypedColumn[UserBean, Double] = udaf.toColumn.name("avgAge")

    // 读取数据，构建 DataFrame
    val df: DataFrame = spark.read.json("in/user.json")
    val ds: Dataset[UserBean] = df.as[UserBean]
    ds.select(avgCol).show()

    spark.stop()
  }
}

case class UserBean(name: String, age: BigInt)
case class AvgBuffer(var sum: BigInt, var count: Int)

/**
  * 声明用户自定义聚合函数(强类型)
  * 1.继承 Aggregator，设定泛型
  * 2.实现方法
  */
class MyAgeAvgClassFunction extends Aggregator[ UserBean, AvgBuffer, Double ]{

  // 初始化
  override def zero: AvgBuffer = AvgBuffer(0, 0)

  // 聚合函数
  override def reduce(b: AvgBuffer, a: UserBean): AvgBuffer = {
    b.sum = b.sum + a.age
    b.count = b.count + 1
    b
  }

  // 缓冲区的合并操作
  override def merge(b1: AvgBuffer, b2: AvgBuffer): AvgBuffer = {
    b1.sum = b1.sum + b2.sum
    b1.count = b1.count + b2.count
    b1
  }

  // 完成计算
  override def finish(reduction: AvgBuffer): Double = reduction.sum.toDouble / reduction.count

  override def bufferEncoder: Encoder[AvgBuffer] = Encoders.product

  override def outputEncoder: Encoder[Double] = Encoders.scalaDouble
}